package com.algorithm.learning.base.图;

import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

/**
 * @program: algorithm-learning
 * @description:
 * @author: YuKai Fan
 * @create: 2024/8/29 22:37
 **/
public class 最小生成树算法之Prim {


    /**
     * 最小生成树的算法之Prim
     * @see 最小生成树算法之Kruskal#kruskalMST(Graph)
     * @param graph
     * @return
     */
    public static Set<Edge> primMST(Graph graph) {
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new EdgeComparator());
        // 哪些点被解锁出来了
        Set<Node> nodeSet = new HashSet<>();
        // 这个其实可以不要，因为 result已经是Set集合，所以即使有重复的边，放到集合里，也不会影响
        // 但是为了防止歧义，这个还是用edgeSet来判断是否有边重复放入
        Set<Edge> edgeSet = new HashSet<>();
        // 依次挑选的边在result中
        Set<Edge> result = new HashSet<>();

        // 要注意，这里可以不用for循环，因为对于P算法，可以从任意一个点出发
        // 这里for循环的目的就是为了表达随机挑一个点
        for (Node node : graph.nodes.values()) {
            // node是开始的点
            if (!nodeSet.contains(node)) {
                nodeSet.add(node);
                // 当前点的所有边，全部入队
                for (Edge edge : node.edges) {
                    if (!edgeSet.contains(edge)) {
                        edgeSet.add(edge);
                        priorityQueue.add(edge);
                    }
                }
                while (!priorityQueue.isEmpty()) {
                    // 出队，这个当前点的最小边
                    Edge edge = priorityQueue.poll();
                    // 这个可能是一个新的点
                    Node toNode = edge.to;
                    // 如果当前toNode不在nodeSet中，就表示这是一个新的点
                    if (!nodeSet.contains(toNode)) {
                        nodeSet.add(toNode);
                        result.add(edge);
                        // 这里很好理解，就是当to点，不在nodeSet中时，就把to点的所有边，全部入队
                        // 那么此时priorityQueue包含了to点的所有的边， 而且是排好序的
                        for (Edge nextEdge : toNode.edges) {
                            if (!edgeSet.contains(nextEdge)) {
                                edgeSet.add(nextEdge);
                                priorityQueue.add(nextEdge);
                            }
                        }
                    }
                }
            }
        }
        return result;
    }


    public static class EdgeComparator implements Comparator<Edge> {

        @Override
        public int compare(Edge o1, Edge o2) {
            return o1.weight - o2.weight;
        }

    }
}
